import torch
import torch.nn as nn
import torch.nn.functional as F
import argparse
import os
import torch.optim as optim
from tensorboardX import SummaryWriter
from util.usher_data_loader import usher_pic_yolo_data_loader
from torch.utils.data import DataLoader
from base_model.yolo.yolov1.yolo_v1_loss import yolo_v1_loss
from datetime import datetime


class train_adapter():
    def __init__(self, model: nn.Module, loss: nn.Module, optimizer: torch.optim.Optimizer = None, model_name: str = "model", device: torch.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")) -> None:
        super().__init__()
        self.model = model
        self.model_name = model_name
        self.loss = loss
        self.optimizer = optimizer
        self.device = device
        self.model.to(self.device)
        self.summary_writer = SummaryWriter(logdir='./log')

    def train(self, train_data, epoch: int):
        # 6.train and record
        size = len(train_data)
        idx = epoch * size
        self.model.train()
        for batch_idx, (data, target) in enumerate(train_data):
            data = data.to(self.device)
            target[0] = target[0].to(self.device)
            target[1] = target[1].to(self.device)
            self.optimizer.zero_grad()

            output = self.model(data)
            loss_sum, (loss, predict_class, ground_class, positive_iou) = self.loss(output, target)

            self.summary_writer.add_scalar(self.model_name, loss_sum.item(), idx)
            idx += 1

            loss_sum.backward()
            self.optimizer.step()
            if batch_idx % 30 == 0:
                print("误差和为:{},\t中心点定位误差:{},\t宽高定位误差:{},\t对的bbox的置信度误差:{},\t对不是bbox的置信度需要和0做误差:{},\t分类误差:{}".format(loss_sum, loss[0], loss[1], loss[2], loss[3], loss[4]))
                print("训练批次:{} loss:{}".format(epoch, loss_sum.item()))

    def test(self, test_data):
        self.model.eval()
        acc = 0
        with torch.no_grad():
            for batch_idx, (data, target) in enumerate(test_data):
                data = data.to(self.device)
                target[0] = target[0].to(self.device)
                target[1] = target[1].to(self.device)

                output = self.model(data)
                loss_sum, (loss, predict_class, ground_class, positive_iou) = self.loss(output, target)
                reshape_iou = positive_iou.reshape(len(positive_iou), 1)
                p = torch.ones(predict_class.shape).to(self.device) - torch.abs(predict_class - ground_class)
                pre_acc = torch.mul(reshape_iou, p) / len(positive_iou)
                acc += pre_acc

        val_accurate = acc / len(test_data)
        print("测试集准确率为:{}".format(val_accurate))


def parse_opt():
    parser = argparse.ArgumentParser()
    parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--restart', default=False, help='是否重新训练,如果重新训练,则使用预训练模型,如果不重新训练,则使用训练好的模型,如果填写对应模型路径,则随机初始化权重')
    parser.add_argument('--per_weight_file', default='', help='官方预训练模型or提前预训练模型的路径 e.g. ../yolo/yolo-c.pt')
    # parser.add_argument('--weight_file', default='', help='上次还没训练完的模型 e.g. ../yolo/my-yolo-c.pt')
    parser.add_argument('--weight_file', default='./model_pt/my-yolo-c.pt', help='上次还没训练完的模型 e.g. ../yolo/my-yolo-c.pt')

    # 获取当前日期和时间
    now = datetime.now()
    # 将日期转换为字符串
    date_str = now.strftime('%Y-%m-%d %H:%M:%S')
    parser.add_argument('--model_name', default='my-yolo-c-' + date_str, help='tensorboard 记录的名字')
    parser.add_argument('--optimizer_name', default='SGD', help='模型优化方法, 支持SGD,Adam,AdamW')
    parser.add_argument('--lr', type=float, default=3e-4, help='初始学习率')
    parser.add_argument('--momentum', type=float, help="动量", default=0.5)
    parser.add_argument('--weight_decay', type=float, help="正则惩罚项系数", default=5e-4)
    parser.add_argument('--epoch_num', type=int, help="训练轮次", default=5)
    parser.add_argument('--classes_num', type=int, help="分类个数,只在训练时存在,且不能和预设权重分类个数不同", default=300)
    parser.add_argument('--batch_size', type=int, help="每训练一个批次,训练的图片数量", default=2)

    return parser.parse_args()


def main(opt):
    # 使用设备
    device = select_device(opt)
    # =========================================================
    # =========================可变项===========================
    # 模型
    from base_model.yolo.yolov1.yolov1 import yolov1
    model = yolov1(classes_num=2).to(device)
    loss = yolo_v1_loss(device=device, classes_num=2)
    # =========================================================
    # 优化器
    optimizer = select_optimizer(model, opt)
    
    # 加载模型参数(如果有的话)
    if opt.restart and os.path.exists(opt.per_weight_file):
        pre_param_dict = torch.load(opt.pre_weight_file, map_location=device)
        model.load_state_dict(pre_param_dict)
        # 手动释放,要不有memory out的风险
        del pre_param_dict
    elif not opt.restart and os.path.exists(opt.weight_file):
        param_dict = torch.load(opt.weight_file, map_location=device)
        model.load_state_dict(param_dict)
        # 手动释放,要不有memory out的风险
        del param_dict

    # 数据
    train_data = usher_pic_yolo_data_loader(2, "F:\\data\\wow\\img_data\\train\\images", "F:\\data\\wow\\img_data\\train\\labels", mean=(0.075, 0.075, 0.071), std=(0.127, 0.127, 0.121))
    test_data = usher_pic_yolo_data_loader(2, "F:\\data\\wow\\img_data\\test\\images", "F:\\data\\wow\\img_data\\test\\labels", mean=(0.075, 0.075, 0.071), std=(0.127, 0.127, 0.121))
    train_loader = DataLoader(train_data, batch_size=opt.batch_size, shuffle=True, num_workers=2, pin_memory=True)
    test_loader = DataLoader(test_data, batch_size=opt.batch_size, shuffle=True, num_workers=2, pin_memory=True)

    # 训练
    train = train_adapter(model, loss, optimizer, model_name=opt.model_name, device=device)
    for epoch in range(opt.epoch_num):
        train.train(train_loader, epoch)
    train.test(test_loader)
    torch.save(model.state_dict(), "./model_pt/my-yolo-c.pt")


def select_optimizer(model, opt):
    if opt.optimizer_name == "SGD":
        return optim.SGD(model.parameters(), lr=opt.lr, weight_decay=opt.weight_decay, momentum=opt.momentum)
    elif opt.optimizer_name == "Adam":
        return optim.Adam(params=model.parameters(), lr=opt.lr, weight_decay=opt.weight_decay)
    elif opt.optimizer_name == "AdamW":
        return optim.AdamW(params=model.parameters(), lr=opt.lr, weight_decay=opt.weight_decay)


def select_device(opt):
    if int(opt.device) in {0, 1, 2, 3}:
        return torch.device("cuda:{}".format(opt.device))
    elif opt.device == "cpu":
        return torch.device("cpu")


if __name__ == '__main__':
    # 1.training parameters
    opt = parse_opt()
    main(opt)
